#!/usr/bin/env python

#Standard ROS imports
import roslib; roslib.load_manifest('cartographer')
import rospy

import sys

#Needed services
from nav_msgs.srv import GetMap
from cartographer.srv import map_srv

#Needed messages
from nav_msgs.msg import OccupancyGrid
from sensor_msgs.msg import LaserScan
from geometry_msgs.msg import PoseWithCovarianceStamped
from geometry_msgs.msg import PoseArray
from geometry_msgs.msg import Quaternion
from cmvision.msg import Blobs

#For ease of use in the updating of the map
from tf.transformations import euler_from_quaternion

#Some math functions
from math import cos
from math import sin
from math import exp
from math import floor
from math import sqrt
from math import pi

#This code assumes that an AMCL node is running, using a pre-made
#map of the area that the robot is going to be navigating around.
#Navigation stack also must be running, so we can retrieve the 
#map from the map_server.
def get_map_client():
    rospy.wait_for_service('static_map')
    try:
        get_map = rospy.ServiceProxy('static_map', GetMap)
        response = get_map()
        return response.map
    except rospy.ServiceException, e:
        rospy.logerr("Service call failed")

#This is for clarity's sake.  This function, given a cell size and number of cells from the mean,
#should return the proper Gaussian value at that cell, in the OccupancyGrid notation of [0,100] in
#probability.
def getGaussian(offX, offY, resolution, variance):
    totalDist = sqrt(pow(offX * resolution, 2) + pow(offY * resolution, 2))
    probability = (1 / sqrt(2 * pi * variance)) * exp(( -pow(totalDist, 2)) / (2 * variance))
    return probability

#This is the primary class that we'll be using.  This creates and maintains a map which dropps 
#Gaussian distributions across a map at locations where it detects people.
class personMap:
    def __init__(self):

        #Sets the variables that will work over all the callbacks.
        rospy.loginfo("Starting the person mapper")

        self.location = PoseWithCovarianceStamped()
        self.localMap = OccupancyGrid()
        self.localMap = get_map_client()

        rospy.loginfo("Map info")
        rospy.loginfo("Resolution: %f", self.localMap.info.resolution)
        rospy.loginfo("Width: %i", self.localMap.info.width)
        rospy.loginfo("Height: %i", self.localMap.info.height)

        #A few things to make our lives easier with manipulating the grid.
        self.dataList = list()
        self.occupied = set()

        #Our publisher, which publishes our current map.
        rospy.loginfo("Beginning the publisher")
        self.pub = rospy.Publisher('personMap', OccupancyGrid)

        #Setting the map to be empty
        rospy.loginfo("Clearing the local map")
        self.clearMap()

        #A few parameters that determine the behavior of the dropped Gaussians.
        #self.personVariance = rospy.get_param('pVariance','1.0')
        #self.decayRate = rospy.get_param('decayRate', '1.01')
        #self.personRadius = rospy.get_param('pRadius', '1.0')
        self.personVariance = 5.0
        self.personRadius = 1.0
        self.decayRate = 1

        #The subscribers for both the localization response and the "person finder", which I currently
        #have as a blob detector.
        rospy.loginfo("Beginning the subscribers")
        rospy.Subscriber('amcl_pose', PoseWithCovarianceStamped, self.locationCallback)
        rospy.Subscriber('blobs', Blobs, self.blobCallback) 

        #For demonstration purposes, here is the line that puts it into laserscan mode.
        #Leave it commented out except for debugging.
        #self.laserUpdate()

        #And finally, loop the degradation function, as well as publishing the map.
        rate = rospy.Rate(1)
        rospy.loginfo("Starting the loop, which publishes and degrades our map")
        while not rospy.is_shutdown():
            rospy.loginfo("Top of loop")
            self.pub.publish(self.localMap)
            self.degrade()
            rate.sleep()


    #updates the location of the robot according to the AMCL node.
    def locationCallback(self, data):
        self.location = data

    #Sets the update to update with a laser scan.
    def laserUpdate(self):
        rospy.Subscriber('base_scan', LaserScan, self.laserCallback)

    #Blob update
    def blobCallback(self, data):
        oneThird = data.image_width/3
        twoThird = data.image_width - oneThird

        for blob in data.blobs:
            if blob.x <= oneThird:
                self.mapUpdate(1.0, -pi/2.0)
            elif blob.x > oneThird and blob.x <= twoThird:
                self.mapUpdate(1.0, 0.0)
            elif blob.x > twoThird and blob.x <= data.image_width:
                self.mapUpdate(1.0, pi/2.0)

    #This is the function that actually updates the map.  It needs the following:
    #Distance from robot
    #Angle from robot
    def mapUpdate(self, dist, angle):

        rospy.loginfo("Map update called with distance: %f and angle: %f", dist, angle)

        #Set local variables to store the width, height, origin, and resolution of the map.
        width = self.localMap.info.width
        height = self.localMap.info.height
        origin = self.localMap.info.origin.position
        resolution = self.localMap.info.resolution

        #Set local variables to store the current position of the robot.
        currX = self.location.pose.pose.position.x
        currY = self.location.pose.pose.position.y
        orientation = self.location.pose.pose.orientation
        robot_angles = euler_from_quaternion([orientation.x, orientation.y, orientation.z, orientation.w], 'sxyz')

        #Calculate the location of the returned target in cartesian coordinates and occGrid cell number.
        objX = currX + (dist * cos(robot_angles[2] + angle))
        objY = currY + (dist * sin(robot_angles[2] + angle))

        cellX = int(floor(objX/resolution) - (origin.x / resolution))
        cellY = int(floor(objY/resolution) - (origin.y / resolution))

        #Small variable to determine how many cells in one direction that our radius for our target covers.
        maxNM = int(floor(self.personRadius/resolution))

        #The actual code that inserts the Gaussian into the map.
        #It's really messy, lots of copy/pastes, but should still work.
        #Need to remember to come back and clean this up.
        for n in range(maxNM):
            for m in range(maxNM):

                #Each if statement ensures that it does not cross any of the boundries.  Each one covers a quadrant of 
                #the gaussian.
                if((cellX - n) >= 0 and (cellY - m) >= 0):
                    self.dataList[int((cellX-n) + ((cellY - m) * width))] += int(100 * getGaussian(n - maxNM / 2, m - maxNM / 2, resolution, self.personVariance))

                    #rospy.loginfo("Q1: Inserting value %i at location %i, %i", int(100 * getGaussian(n - maxNM / 2, m - maxNM / 2, resolution, self.personVariance)), cellX, cellY)

                    newData = int((cellX - n) + (cellY - m) * width)

                    if self.dataList[newData] > 100:
                        self.dataList[newData] = 100

                    if newData not in self.occupied:
                        self.occupied.add(newData)

                if((cellX - n) >= 0 and (cellY + m) < height):
                    self.dataList[int((cellX-n) + ((cellY + m) * width))] += int(100 * getGaussian(n - maxNM / 2, m - maxNM / 2, resolution, self.personVariance))
                    
                    #rospy.loginfo("Q2: Inserting value %i at location %i, %i", int(100 * getGaussian(n - maxNM / 2, m - maxNM / 2, resolution, self.personVariance)), cellX, cellY)

                    newData = int((cellX - n) + (cellY + m) * width)

                    if self.dataList[newData] > 100:
                        self.dataList[newData] = 100

                    if newData not in self.occupied:
                        self.occupied.add(newData)

                if((cellX + n) < width and (cellY - m) >= 0):
                    self.dataList[int((cellX+n) + ((cellY - m) * width))] += int(100 * getGaussian(n - maxNM /2, m - maxNM / 2, resolution, self.personVariance))

                    #rospy.loginfo("Q3: Inserting value %i at location %i, %i", int(100 * getGaussian(n - maxNM / 2, m - maxNM / 2, resolution, self.personVariance)), cellX, cellY)

                    newData = int((cellX + n) + (cellY - m) * width)

                    if self.dataList[newData] > 100:
                        self.dataList[newData] = 100

                    if newData not in self.occupied:
                        self.occupied.add(newData)

                if((cellX + n) < width and (cellY + m) < height):
                    self.dataList[int((cellX+n) + ((cellY + m) * width))] += int(100 * getGaussian(n - maxNM / 2, m - maxNM / 2, resolution, self.personVariance))

                    #rospy.loginfo("Q4: Inserting value %i at location %i, %i", int(100 * getGaussian(n - maxNM / 2, m - maxNM / 2, resolution, self.personVariance)), cellX, cellY)

                    newData = int((cellX + n) + (cellY + m) * width)

                    if self.dataList[newData] > 100:
                        self.dataList[newData] = 100

                    if newData not in self.occupied:
                        self.occupied.add(newData)

        #And then change out the map for the new one.
        self.localMap.data = tuple(self.dataList)




    #This is using a laserscan, and is a bit deprecrated.
    #WARNING: using this callback will make it run incredibly slow.  Possibly even overlap
    #the callbacks, due to the speed of the laser scan.  THIS IS NOT RECOMMENDED.
    def laserCallback(self, data):
        count = 0

        for each_item in data.ranges:
            theta = (count * data.angle_increment) + data.angle_min
            self.mapUpdate(each_item, theta)
            count += 1

    #Degrades our occupancy grid over time.
    def degrade(self):
        tempSet = set()

        while len(self.occupied) > 0:
            index = self.occupied.pop()
            if self.dataList[index] > 0:
                self.dataList[index] = self.dataList[index] - 5
                if self.dataList[index] < 0:
                    self.dataList[index] = 0

            if self.dataList[index] > 0:
                tempSet.add(index)

        self.localMap.data = tuple(self.dataList)
        self.occupied ^= tempSet

    #Resets the map to an empty map.
    def clearMap(self):
        for i in range(self.localMap.info.width * self.localMap.info.height):
            self.dataList.append(0)

        self.localMap.data = tuple(self.dataList)

#The main function.
if __name__ == '__main__':
    rospy.init_node('personMapper')

    #Creates the map, then spins until the node is killed.
    mapper = personMap()

    #And it spins until the node is shut down.
    rospy.spin()
